package com.duqc.util;

import io.lettuce.core.output.KeyValueStreamingChannel;
import org.apache.commons.collections4.CollectionUtils;

import java.util.*;
import java.util.function.Function;

/**
 * @ClassName TreeNode
 * @Description TreeNode
 * @Author duqc
 * @Date 2023/4/26
 **/
public class TreeNode<T> {

    private final T entity;

    private final List<TreeNode<T>> children;

    public TreeNode(T entity) {
        this.entity = entity;
        this.children = new ArrayList<>();
    }

    public T getEntity() {
        return entity;
    }

    public List<TreeNode<T>> getChildren() {
        return children;
    }

    public void addChild(TreeNode<T> child) {
        children.add(child);
    }

    /**
     * 构建树状结构
     *
     * @param entities    原始数据
     * @param id          数据ID
     * @param pid         数据父级ID
     * @param setChildren 下级数据
     * @return List<T>
     */
    public static <T, K> List<T> build(List<T> entities, Function<T, K> id, Function<T, K> pid, KeyValueStreamingChannel<T, List<T>> setChildren, String... rootId) {
        Map<K, TreeNode<T>> nodeMap = new HashMap<>(entities.size());
        List<String> rootIdList = Arrays.asList(rootId);

        for (T entity : entities) {
            K entityId = id.apply(entity);
            nodeMap.put(entityId, new TreeNode<>(entity));
        }

        List<TreeNode<T>> roots = new ArrayList<>();
        for (T entity : entities) {
            K entityId = id.apply(entity);
            K parentEntityId = pid.apply(entity);
            TreeNode<T> node = nodeMap.get(entityId);
            TreeNode<T> parentNode = nodeMap.get(parentEntityId);
            if (null != parentNode) {
                parentNode.addChild(node);
            }
            if (null == parentNode && CollectionUtils.isNotEmpty(rootIdList) && rootIdList.contains(String.valueOf(entityId))) {
                roots.add(node);
            }
            if (null == parentNode && CollectionUtils.isEmpty(rootIdList)) {
                roots.add(node);
            }
        }
        return list(roots, setChildren);
    }

    public static <T> List<T> list(List<TreeNode<T>> entities, KeyValueStreamingChannel<T, List<T>> setChildren) {
        List<T> list = new ArrayList<>();
        if (CollectionUtils.isNotEmpty(entities)) {
            entities.forEach(treeNode -> {
                T entity = treeNode.getEntity();
                list.add(entity);
                List<T> listChildren = list(treeNode.getChildren(), setChildren);
                setChildren.onKeyValue(entity, listChildren);
            });
        }
        return list;
    }

}
